package com.aispeech.asr.apis.ssc.beans;

import com.aispeech.asr.apis.ssc.config.AIOpsConfig;
import com.aispeech.asr.apis.ssc.config.AppConfig;
import com.aispeech.asr.apis.ssc.data.domain.SscContext;
import com.aispeech.asr.apis.ssc.data.form.Audio;
import com.aispeech.asr.apis.ssc.data.form.SscForm;
import com.aispeech.asr.clib.arpc.*;
import com.aispeech.asr.clib.decoder.EAudioCodec;
import com.aispeech.asr.clib.utils.SplittedPcmData;
import com.aispeech.asr.clib.vad.EVFrameState;
import com.aispeech.asr.comm.exceptions.AsrException;
import com.fasterxml.jackson.databind.JsonNode;
import com.fasterxml.jackson.databind.ObjectMapper;
import com.sun.jna.Pointer;
import lombok.SneakyThrows;
import lombok.extern.slf4j.Slf4j;
import org.springframework.beans.factory.annotation.Autowired;
import org.springframework.beans.factory.config.ConfigurableBeanFactory;
import org.springframework.context.annotation.Scope;
import org.springframework.stereotype.Component;

import java.io.FileInputStream;
import java.io.IOException;
import java.io.InputStream;
import java.nio.charset.StandardCharsets;
import java.nio.file.Files;
import java.nio.file.Path;
import java.nio.file.Paths;
import java.util.concurrent.atomic.AtomicInteger;
import java.util.concurrent.atomic.AtomicReference;
import java.util.function.Function;

@Component
@Slf4j
@Scope(ConfigurableBeanFactory.SCOPE_PROTOTYPE)
public class Vad implements Function<SscContext, SscContext>, IARpcCallback {
    @Autowired
    AppConfig config;
    @Autowired
    AIOpsConfig aiOpsConfig;
    @Autowired
    ObjectMapper objectMapper;

    private AsrException exception;
    private AtomicReference<EVFrameState> evFrameState;
    private AtomicInteger lastBeginOffset;
    private SplittedPcmData splittedPcmData;
    private SscContext context;

    private Object waiter = new Object();
    private ARpc rpc;

    public void callback(String sessionid, int state, Pointer data, int size, Pointer userData) {
//        log.info("[ Vad.callback ] session: {}, state: {}, size: {}", sessionid, state, size);

        if ((state == EARpcState.S_CONNECTION_BROKEN.getValue()) || (state == EARpcState.S_CONNECTION_CLOSED.getValue())) {
            exception = new AsrException(String.format("[ Vad.callback ] session: %s, state: %d", sessionid, state));
        } else if (state == EARpcState.S_INTERNAL_ERROR.getValue()) {
            byte[] bytes = data == null ? new byte[0] : data.getByteArray(0, size);
            exception = new AsrException(String.format("[ Vad.callback ] session: %s, state: %d, reason: %s",
                    sessionid, state, new String(bytes, StandardCharsets.UTF_8)));
        } else {
            if (size > 0) {
                try {
                    JsonBin jsonBin = ARpc.parseJsonBinByStream(data, size);

                    if (jsonBin != null) {
                        if (jsonBin.getBin() != null)
                            splittedPcmData.write(jsonBin.getBin());

                        JsonNode res = objectMapper.readTree(jsonBin.getJson());
                        int code = res.get("code").asInt();

                        if (code == 0) {
                            JsonNode jsonNode = res.get("state_offset_list");
                            if (jsonNode != null && jsonNode.size() > 0) {
                                for (int i = 0; i < jsonNode.size(); i++) {
                                    JsonNode item = jsonNode.get(i);
                                    int vadState = item.get("state").asInt();
                                    int offset = item.get("offset").asInt();

                                    evFrameState.set(EVFrameState.from(vadState));

                                    if (evFrameState.get() == EVFrameState.SIL) {
                                        splittedPcmData.addOffset(lastBeginOffset.get(), offset);
                                    } else if (evFrameState.get() == EVFrameState.SPEECH_BEGIN) {
                                        lastBeginOffset.set(offset);
                                    }
                                }
                            }
                        } else {
                            exception = new AsrException(String.format("vad failed. code: %s, bytes: %s",
                                    code, res.get("message")));
                        }
                    } else {
                        exception = new AsrException(String.format("parse vad result failed. session: %s, bytes: %s",
                                sessionid, data.getByteArray(0, size)));
                    }
                } catch (Exception e) {
                    e.printStackTrace();
                    exception = new AsrException(String.format("parse vad result failed. session: %s, bytes: %s, error: %s",
                            sessionid, data.getByteArray(0, size), e.getMessage()));
                }
            } else {
                exception = new AsrException(String.format("[ Vad.callback ] session: %s, state: %d, size length is zero", sessionid, state));
            }
        }

        synchronized (waiter) {
            waiter.notifyAll();
        }
    }

    @SneakyThrows
    @Override
    public SscContext apply(SscContext context) {
        context.addMetric("vad_begin");

        this.context = context;
        this.exception = null;

        this.splittedPcmData = new SplittedPcmData();
        this.evFrameState = new AtomicReference<>(EVFrameState.SIL);
        this.lastBeginOffset = new AtomicInteger();

        rpc = new ARpc();
        try {
            prepareAndStartVad();
            doVad();
            context.setOffsetList(splittedPcmData.getOffsetList());
            savePcmDataToWorkDir();

            splittedPcmData = null;

            context.addMetric("vad_end");
        } finally {
            rpc.close();
            rpc = null;
        }

        return context;
    }

    @SneakyThrows
    private void savePcmDataToWorkDir() {
        SscForm sscForm = context.getForm();
        Audio audio = sscForm.getAudio();
        Path pcmPath = Paths.get(context.getAudioFile().getParent(), context.getSessionId() + ".pcm");

        byte[] pcm = splittedPcmData.getBuffer().array();
        Files.write(pcmPath, pcm);
        context.setPcmPath(pcmPath);

        sscForm.setDuration(
                AppConfig.pcmBytesToMsecs(
                        pcm.length,
                        audio.getSampleRate(),
                        audio.getChannel(),
                        audio.getSampleBytes())
        );
    }

    private void doVad() throws IOException, InterruptedException, AsrException {
        String sessionId = rpc.sessionId();
//        String audioType = context.getForm().getAudio().getAudioType();
        int size = 128000;
        int timeout = 5 * 60;
        int len;
//        if (!(audioType.equalsIgnoreCase("pcm") || audioType.equalsIgnoreCase("wav"))) {
//            size /= 8;
//        }

        InputStream inputStream = new FileInputStream(context.getAudioFile());
        byte[] buff = new byte[size];

        while ((len = inputStream.read(buff)) > 0) {
            rpc.feed(buff, len);
//            log.info("[ doVad ]  sessionid: {}, feed length: {}...", sessionId, len);
            synchronized (waiter) {
                waiter.wait(100);
            }

            if (exception != null)
                throw exception;
        }
        inputStream.close();

        ArpcOutput output = rpc.stop();
//        if (output.getCode() == EARpcState.S_OK.getValue()) {
        if (output.getCode() == EARpcState.S_EOF.getValue()) {
            byte[] data = output.getData();
            if (data != null && data.length > 0) {
                JsonBin jsonBin = ARpc.parseJsonBinByBytes(data);

                if (jsonBin != null) {
                    if (jsonBin.getBin() != null)
                        splittedPcmData.write(jsonBin.getBin());

                    JsonNode res = objectMapper.readTree(jsonBin.getJson());
                    int code = res.get("code").asInt();

                    if (code == 0) {
                        JsonNode jsonNode = res.get("state_offset_list");
                        if (jsonNode != null && jsonNode.size() > 0) {
                            for (int i = 0; i < jsonNode.size(); i++) {
                                JsonNode item = jsonNode.get(i);
                                int vadState = item.get("state").asInt();
                                int offset = item.get("offset").asInt();

                                evFrameState.set(EVFrameState.from(vadState));

                                if (evFrameState.get() == EVFrameState.SIL) {
                                    splittedPcmData.addOffset(lastBeginOffset.get(), offset);
                                } else if (evFrameState.get() == EVFrameState.SPEECH_BEGIN) {
                                    lastBeginOffset.set(offset);
                                }
                            }
                        }
                    } else {
                        throw new AsrException(String.format("[ doVad ] failed stop rpc. session: {}, code: %s, bytes: %s",
                                rpc.sessionId(), code, res.get("message")));
                    }
                } else {
                    throw new AsrException(String.format("[ doVad ] failed stop rpc. session: {}, failed parse last data: {}", rpc.sessionId(), data));
                }
            }
        } else {
            throw new AsrException(String.format("[ doVad ] failed stop rpc. session: {}, code: {}", rpc.sessionId(), output.getCode()));
        }


        log.info("[ doVad.finished ] sessionId: {}, vad slices: {}", sessionId,
                splittedPcmData != null ? splittedPcmData.getOffsetList().size() : -1);
    }

    private void prepareAndStartVad() throws AsrException {
        int destSampleRate;
        String vadRes;

        Audio audio = context.getForm().getAudio();
        EAudioCodec audioCodec = EAudioCodec.from(audio.getAudioType());

        if (audioCodec == null) throw new AsrException("Unknwon audio type: " + audio.getAudioType());

        if (audio.getSampleRate() == 8000) {
            destSampleRate = 8000;
            vadRes = String.format("%s/vad/%s", "", config.getVadRes8k());
        } else {
            destSampleRate = 16000;
            vadRes = String.format("%s/vad/%s", "", config.getVadRes16k());
        }

        String param_decoder = String.format("{\"dst_rate\":%d,\"audio_type\":%d,\"src_rate\":%d,\"channel\":%d,\"samplebytes\":%d}",
                destSampleRate, audioCodec.value(), audio.getSampleRate(), audio.getChannel(), audio.getSampleBytes());
        String param_vad = String.format("{\"cfg_path\":\"%s\",\"is_bin\":1,\"pause_time\":0}", vadRes);

        String params = String.format("{\"op\":\"vad\",\"audio-decoder\":%s,\"vad\":%s}", param_decoder, param_vad);

        rpc.start(aiOpsConfig.getHost(), aiOpsConfig.getPort(), this, params, 3);
        log.info("[ doPrepareVad ] sessionid: {}, params: {}", rpc.sessionId(), params);
    }
}
